#ifndef YOLT_H
#define YOLT_H

#include "./tiny_dnn/tiny_dnn.h"
#include "yolt_def.hpp"
#include <opencv2/opencv.hpp>

namespace nn_yolt
{

class NNYolt
{

    typedef tiny_dnn::network<tiny_dnn::sequential> nn_t;
    typedef tiny_dnn::convolutional_layer C;
    typedef tiny_dnn::max_pooling_layer S;
    typedef tiny_dnn::fully_connected_layer F;

    typedef tiny_dnn::batch_normalization_layer B;

    typedef tiny_dnn::activation::leaky_relu leaky_relu;

    typedef tiny_dnn::activation::relu relu;

    typedef tiny_dnn::vec_t f_vec_t;

  public:
    NNYolt(const int _in_width = 416, const int _in_height = 416,
           const int _in_channels = 3, const int _out_cell_num_hor = 26,
           const int _out_cell_num_ver = 26, const int _out_box_num_per_cell = 2,
           const str_t & net_file_name = "net.dat")
        : in_width(_in_width), in_height(_in_height), in_channels(_in_channels),
          out_width(_out_cell_num_hor), out_height(_out_cell_num_ver),
          box_num_per_cell(_out_box_num_per_cell),
          out_depth(_out_box_num_per_cell * PARAM_NUM_PER_BOX),
          net_file_name(net_file_name)
    {
        ;
    }

    ~NNYolt() { ; }

    void construct_net();

    void save_net(const str_t &saved_name)
    {
        net.save(saved_name, tiny_dnn::content_type::weights_and_model, tiny_dnn::file_format::json);
    }

    void load_net(const std::string &filepath)
    {
        net.load(filepath, tiny_dnn::content_type::weights_and_model, tiny_dnn::file_format::json);
    }

    void gen_connection_table(const int in_size, const int out_size, bool *connection_table);

    void train(const str_t samples_dir, const int train_index_begin,
               const int train_sample_num, const int batch_size, const int epoch);
    void test(const str_t samples_dir, const int test_index_begin, const int test_num);

  private:
    void mat_2_in_data(const str_t image_path, f_vec_t &vec) const;
    void label_2_out_data(const YoltLabel &label, f_vec_t &vec) const;

    void in_data_2_mat(const f_vec_t &in_data, cv::Mat &mat) const;
    void out_data_2_label(const f_vec_t &out_data, YoltLabel &label) const;

    void get_train_data(const str_t samples_dir,
                        std::vector<f_vec_t> &in_datas,
                        std::vector<f_vec_t> &out_datas,
                        const int train_sample_index_begin = 0,
                        const int train_sample_num = -1) const;

  private:
    int in_width, in_height, in_channels;
    int out_width, out_height, out_depth;
    int box_num_per_cell;
    nn_t net;
    const str_t net_file_name;
};

} // namespace nn_yolt

#endif //YOLT_H
